package render;

import geom.Geometry;

import java.awt.Color;
import java.awt.Component;
import java.awt.Graphics2D;
import java.awt.RenderingHints;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferInt;
import java.io.File;
import java.io.IOException;
import java.text.DecimalFormat;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.Properties;
import java.util.Random;

import javax.imageio.ImageIO;

import main.TracerProperties;
import util.KDNode;

/**
 * buffered ray-tracing renderer 
 * 
 * @author Robbe
 *
 */
public class RaytracerImage extends BufferedImage {
	
	private static int[][] distributions = {
		{
			0,0,1,0,0,
			0,1,1,1,0,
			1,1,1,1,1,
			0,1,1,1,0,
			0,0,1,0,0
		},			
		{
			0,0,0,1,0,0,0,
			0,0,1,1,1,0,0,
			0,1,1,1,1,0,0,
			1,1,1,1,1,1,1,
			0,1,1,1,1,1,0,
			0,0,1,1,1,0,0,
			0,0,0,1,0,0,0
		},
		{
			0,0,0,0,1,0,0,0,0,
			0,0,1,1,1,1,1,0,0,
			0,1,1,1,1,1,1,1,0,
			0,1,1,1,1,1,1,0,0,
			1,1,1,1,1,1,1,1,1,
			0,1,1,1,1,1,1,1,0,
			0,0,1,1,1,1,1,1,0,
			0,0,1,1,1,1,1,0,0,
			0,0,0,0,1,0,0,0,0,
		},		
	};
	
	private int[] buffer; // direct manipulation using frame buffer
	private Graphics2D g2d; // manipulation using java renderer	
	
	private static final int clearColor = Color.GRAY.getRGB();
	
	public static final double tMin = 0.01;
	
	// status update	
	private Component observer;
	private double oldPercent;
	
	// statistics
	private long intersectionChecks;
	private long raysTraced;
	private long timeStart;
	
	// properties
	private Properties properties;
	private boolean shadowProbe;

	private boolean reflections;
	private boolean transmissions;
	
	private boolean diffuseShading;
	private boolean specularShading;	
	private boolean ambientShading;
	
	private boolean whittedSampling;
	private boolean dof;
	private double samplingFraction;
	private boolean diffuse;
	private int diffuseDensity;
	private int maxLevel;
	
	private boolean softShadows;
	private int softShadowDensity;
	
	RaytracerImage parent;
	
	KDNode tree;
	
	Random rand;

	public RaytracerImage(Properties props, Component observer) {
		// create image buffer
		super(Integer.parseInt(props.getProperty("output_width","640")), Integer.parseInt(props.getProperty("output_height","480")), BufferedImage.TYPE_INT_RGB);
		
		this.observer = observer;
		parent = null;
		
		// extract properties
		this.properties = props != null ? props : new TracerProperties();
		diffuseShading = Boolean.parseBoolean(properties.getProperty("shading_diffuse"));
		specularShading = Boolean.parseBoolean(properties.getProperty("shading_specular"));
		ambientShading = Boolean.parseBoolean(properties.getProperty("shading_ambient"));		
		shadowProbe = Boolean.parseBoolean(properties.getProperty("trace_shadows"));
		reflections = Boolean.parseBoolean(properties.getProperty("trace_reflections"));
		transmissions = Boolean.parseBoolean(properties.getProperty("trace_transmissions"));
		whittedSampling = Boolean.parseBoolean(properties.getProperty("sampling_whitted"));
		dof = Boolean.parseBoolean(properties.getProperty("dof"));
		samplingFraction = Double.parseDouble(properties.getProperty("sampling_fraction","1.0"));
		diffuse = Boolean.parseBoolean(properties.getProperty("diffuse"));
		diffuseDensity = Integer.parseInt(properties.getProperty("diffuse_density"));
		maxLevel = Integer.parseInt(properties.getProperty("trace_level"));
		
		softShadows = Boolean.parseBoolean(properties.getProperty("softShadows","false"));
		softShadowDensity = Integer.parseInt(properties.getProperty("softShadows_density","10"));
		
		// initialize frame buffer manipulators
		buffer = ((DataBufferInt)getRaster().getDataBuffer()).getData();
		g2d = (Graphics2D) getGraphics();
		
		// set some rendering hints (let sampled images look better)
		RenderingHints renderingHints = new RenderingHints(null);
		renderingHints.put(RenderingHints.KEY_RENDERING, RenderingHints.VALUE_RENDER_QUALITY);
		renderingHints.put(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
		g2d.addRenderingHints(renderingHints);		
		
		// draw wait screen
		clear();
		g2d.setColor(Color.BLACK);
		g2d.drawString("please wait ...", getWidth()/2, getHeight()/2);
		
		rand = new Random();
	}
	
	public void clear() {
		Arrays.fill(buffer, clearColor);
	}
	
	private void statusUpdate(double percent) {		
		if (oldPercent + 1 < percent) {
			System.out.println("progress : " + (int)percent+ "%");
			oldPercent = percent;
			
			if (parent != null) {
				//g2d.drawImage(sampled, 0, 0, null);
				parent.g2d.drawImage(this, 0, 0, null);
			}
			
			if (observer != null) {
				observer.repaint();
			}			
		}		
	}
	
	public void draw(Scene scene) {

		// build tree		
		LinkedList<Primitive> primitves = new LinkedList<Primitive>();
		for (int i = 0; i < scene.primitives.length; i++) {
			primitves.add(scene.primitives[i]);
		}
		tree = KDNode.createKDTree(primitves, 0);
		tree.printLeafs();
		
		// reset statistics
		intersectionChecks = 0;
		raysTraced = 0;
		timeStart = System.currentTimeMillis();
		oldPercent = 0;
		
		boolean motionBlur = Boolean.parseBoolean(properties.getProperty("motionBlur","false"));
		double increment = Double.parseDouble(properties.getProperty("motionBlur_increment","0.1"));
		int steps = Integer.parseInt(properties.getProperty("motionBlur_steps","10"));
		
		if (samplingFraction != 1) {
			drawSampled(scene, samplingFraction);
			return;
		}
		
		// choose desired rendering method
		if (whittedSampling) {
			drawWhitted(scene);
		} else if (dof) {
			drawDOF(scene);
		} else if (motionBlur) {
			drawMotion(scene, increment, steps);
		} else {
			drawSingleRay(scene);
		}
		
		//drawTime(scene, 10);
		
		// print some statistics
		final long msPerSecond = 1000;
		final long msPerMinute = msPerSecond * 60;
		final long msPerHour = msPerMinute * 60;			
		
		long time = System.currentTimeMillis() - timeStart;
		StringBuilder sbTime = new StringBuilder();
		sbTime.append(time / msPerHour).append(":").
		append(new DecimalFormat("00").format(time / msPerMinute)).append(":").
		append(new DecimalFormat("00").format(time / msPerSecond)).append("'").
		append(new DecimalFormat("000").format(time % msPerSecond));		
		
		System.out.println("Intersection checks : " + intersectionChecks);
		System.out.println("Traced rays : " + raysTraced);		
		System.out.println("Render time : " + sbTime.toString());
	}	

	private void drawSingleRay(Scene scene) {
		Camera camera = scene.camera;
		
		double[] color = {0,0,0};
		double[] ray;
		
		int width = getWidth();
		int height = getHeight();		
		
		for (int x = 0; x < width; x++) {
			for (int y = 0; y < height; y++) {
				// one ray trough pixel center
				ray = camera.getRay( (x+0.5)/getWidth() - 0.5, (y+0.5)/getHeight() - 0.5);
				
				// trace ray
				color = trace(scene, camera.eye, ray, null, 1);
				
				// map color to RGB integer (8 bit each channel) and write to frame buffer
				buffer[y*width + x] = mapColor(color, 1);				
			}
			
			// refresh view
			statusUpdate((100.0 * x) / width);
		}
	}
	
	private void drawMotion(Scene scene, double increment, int steps) {
		Camera camera = scene.camera;
		
		double[] color = {0,0,0};
		double[] ray;
		
		int width = getWidth();
		int height = getHeight();		
		
		for (int x = 0; x < width; x++) {
			for (int y = 0; y < height; y++) {
				color = new double[]{0,0,0};
				
				// trace rays (camera displacement each frame)
				Camera c = new Camera(camera.eye,camera.gaze,camera.up,camera.fov,camera.ratio);
				for (int i = 0; i < steps; i++) {
					c = new Camera(Geometry.sum(c.eye,new double[]{0,0,-increment}), c.gaze, c.up, c.fov, c.ratio);
					ray = c.getRay( (x+0.5)/getWidth() - 0.5, (y+0.5)/getHeight() - 0.5);
					Geometry.add(color, trace(scene, c.eye, ray, null, 1));
				}
				Geometry.multiply(color, 1.0/steps);
				
				// map color to RGB integer (8 bit each channel) and write to frame buffer
				buffer[y*width + x] = mapColor(color, 1);				
			}
			
			// refresh view
			statusUpdate((100.0 * x) / width);
		}
	}	
	
	private void drawDOF(Scene scene) {
		// read properties
		final int distId = Integer.parseInt(properties.getProperty("dof_distribution","0"));
		final double focalDistance = Double.parseDouble(properties.getProperty("dof_focalDistance","1.0"));
		final double lensSize = Double.parseDouble(properties.getProperty("dof_lensSize","0.1"));;
		final double scatter = Double.parseDouble(properties.getProperty("dof_scatter","0.001"));;
		
		// distribution 
		int[] distribution = distributions[distId];
		int distSize = (int)Math.sqrt(distribution.length);
		double distX, distY;
		double distDelta = lensSize / distSize;
		
		Camera camera = scene.camera;
		
		double[] color = {0,0,0};
		double[] focus, start, ray;
		
		int width = getWidth();
		int height = getHeight();
		
		int numOfRays;
		
		Random rand = new Random(); // used for scatter
		
		// count number of rays in distribution
		numOfRays = 0;
		for (int i = 0; i < distribution.length; i++) {
			if (distribution[i] != 0) {
				numOfRays++;
			}
		}
		
		// trace all view rays
		for (int x = 0; x < width; x++) {
			for (int y = 0; y < height; y++) {
				// main ray
				ray = camera.getRay( (x+0.5)/getWidth() - 0.5, (y+0.5)/getHeight() - 0.5);				
				focus = Geometry.linearCombination(1, camera.eye, focalDistance, ray);
				
				// trace and integrate distribution
				for (int i = 0; i < distribution.length; i++) {
					if (distribution[i] != 0) {
						distX = (i % distSize) - distSize/2.0;
						distY = (i / distSize) - distSize/2.0;
						
						start = Geometry.linearCombination(1,camera.eye, (rand.nextDouble() % scatter) + distX*distDelta, camera.left, (rand.nextDouble() % scatter) + distY*distDelta, camera.up);
						ray = Geometry.difference(focus, start);
						Geometry.normalize(ray);
						
						Geometry.add(color, trace(scene, start, ray, null, 1));
					}
				}
				Geometry.multiply(color, 1.0 / numOfRays);
				
				// map color to RGB integer (8 bit each channel) and write to frame buffer
				buffer[y*width + x] = mapColor(color, 1);				
			}
			
			// refresh view
			statusUpdate((100.0 * x) / width);
		}		
	}
	
	// TODO enable whitted sampling + DOF
	public void drawWhitted(Scene scene) {
		Camera camera = scene.camera;
		
		double[] color = {0,0,0};
		double[] ray;
		
		double[][] whittedBuffer = new double[(getWidth()+1) * (getHeight()+1)][];
		
		int width = getWidth();
		int height = getHeight();
		int whittedWidth = width + 1;
		
		double u,v;
		final double[] zeroVector = {0,0,0};
		for (int i = 0; i < whittedBuffer.length; i++) {
			whittedBuffer[i] = zeroVector;
		}
		
		// build oversampled buffer
		for (int x = 0; x <= getWidth(); x++) {
			for (int y = 0; y <= getHeight(); y++) {
				// one ray through raster crossings (whitted)
				u = ((double)x)/width- 0.5;
				v = ((double)y)/height - 0.5;				
				ray = camera.getRay(u, v);
				
				// trace ray
				color = trace(scene, camera.eye, ray, null, 1);
				
				// write data to whitted buffer
				whittedBuffer[y*whittedWidth + x] = color;
				
				// average data to image buffer
				if (x > 0 && y > 0) {
					color[0] = (whittedBuffer[(y)*whittedWidth+(x)][0] + whittedBuffer[(y)*whittedWidth+(x-1)][0] + whittedBuffer[(y-1)*whittedWidth+(x)][0] + whittedBuffer[(y-1)*whittedWidth+(x-1)][0]) / 4;
					color[1] = (whittedBuffer[(y)*whittedWidth+(x)][1] + whittedBuffer[(y)*whittedWidth+(x-1)][1] + whittedBuffer[(y-1)*whittedWidth+(x)][1] + whittedBuffer[(y-1)*whittedWidth+(x-1)][1]) / 4;
					color[2] = (whittedBuffer[(y)*whittedWidth+(x)][2] + whittedBuffer[(y)*whittedWidth+(x-1)][2] + whittedBuffer[(y-1)*whittedWidth+(x)][2] + whittedBuffer[(y-1)*whittedWidth+(x-1)][2]) / 4;
					
					buffer[(y-1)*width+(x-1)] = mapColor(color, 1);
				}
				
				// refresh view
				statusUpdate((100.0 * x) / whittedWidth);				
			}
		}
	}
	
	public void drawSampled(Scene scene, double fraction) {
		final int width = Integer.parseInt(properties.getProperty("output_width","640"));
		final int height = Integer.parseInt(properties.getProperty("output_height","480"));
		final int sampledWidth = (int)(width * fraction);
		final int sampledHeight = (int)(height * fraction);
		final double scale = (double) width / sampledWidth;
		
		// properties for sampled rendering
		Properties sampledProperties = (Properties)properties.clone();
		sampledProperties.setProperty("output_width", Integer.toString(sampledWidth));
		sampledProperties.setProperty("output_height", Integer.toString(sampledHeight));
		sampledProperties.setProperty("sampling_fraction", Double.toString(1.0));
		
		RaytracerImage sampled = new RaytracerImage(sampledProperties, observer);
		sampled.parent = this;
		
		g2d.scale(scale, scale);
		sampled.draw(scene);
		g2d.scale(1/scale, 1/scale);
	}
	
	private int mapColor(double[] intensity, double gamma) {
		// map arbitrary intensity values to RGB integers  
		return (gammaCorrection(intensity[2], gamma)) | (gammaCorrection(intensity[1], gamma)<<8) | (gammaCorrection(intensity[0], gamma)<<16);
	}
	
	private int gammaCorrection(double intensity, double gamma) {
		// map intensity to [0,1]
		if (intensity > 1) {
			intensity = 1;
		} else if (intensity < 0) {
			intensity = 0;
		}
		
		// gamma correction
		intensity = Math.pow(intensity, gamma);
		
		// scale to size of one byte and round
		return (int) (255 * intensity);
	}
	
	private Primitive intersectionKD(Scene scene, double[] pos, double[] ray, double length, double[] intersection) {		
		//KDNode kdTree = scene.kdTree;
		return tree.travers(pos, ray, length, intersection);
	}
	
	public static Primitive intersection(Scene scene, Primitive exlude, double[] pos, double[] ray, double length, double[] intersection) {		
		Primitive[] primitives = scene.primitives;
		double t, tClosest = Double.MAX_VALUE;
		int iClosest = -1;
		
		//intersectionChecks++;
		
		// find first intersection of ray with a primitive in the scene
		for (int i = 0; i < primitives.length; i++) {
			if (primitives[i] != exlude) { // skip excluded primitive
				t = primitives[i].intersection(pos, ray);
				if (t != Double.NaN && t >= tMin && t < tClosest && t <= length) { // intersecting and on positive side of ray and closer than closest primitive so far				
					tClosest = t;
					iClosest = i;
				}
			}
		}
		
		if (iClosest > -1) { // intersection found
			if (intersection != null) {
				intersection[0] = pos[0] + ray[0] * tClosest;
				intersection[1] = pos[1] + ray[1] * tClosest;
				intersection[2] = pos[2] + ray[2] * tClosest;
			}
			return primitives[iClosest];
		} else {
			return null;
		}
	}
	
	private double[] trace(Scene scene, double[] pos, double[] ray, Primitive exlude, int level) {
		double[] color = {0,0,0};		
		Primitive intersected;
		
		double[] intersection = {0,0,0};
		double[] reflection;
		double[] transmission = {0,0,0};
		double[] normal = null;
		
		raysTraced++;
		
		if (level <= maxLevel) {
			intersected = intersection(scene, exlude, pos, ray, Double.MAX_VALUE, intersection);
			//intersected = intersectionKD(scene, pos, ray, Double.MAX_VALUE, intersection);
			
			if (intersected != null) {
				normal = intersected.getNormal(intersection);
				reflection = Geometry.sum(Geometry.product(normal, -2*Geometry.dotProduct(ray, normal)), ray);
				Geometry.add(color, shade(scene, intersected, intersection, normal, reflection));				
				
				if (reflections) {					
					if (intersected.getMaterial().reflection != 0) {
						Geometry.add(color, Geometry.product(trace(scene, intersection, reflection, intersected, level + 1), intersected.getMaterial().reflection));
					}
				}
				
				if (transmissions) {
					if (intersected.getMaterial().translucence != 0) {
						// TODO correct transmission
						transmission = ray;
						Geometry.add(color, Geometry.product(trace(scene, intersection, transmission, intersected, level + 1), intersected.getMaterial().translucence));					
					}
				}
				
				if (diffuse) {
					Random rand = new Random();
					double[] diff = {0,0,0};
					double[] random, direction, diffuseRay;
					diffuse = false;
					for (int i = 0; i < diffuseDensity; i++) {
						random = new double[]{rand.nextDouble()%1,rand.nextDouble()%1,rand.nextDouble()%1};
						direction = Geometry.crossProduct(normal, random);
						Geometry.normalize(direction);						
						
						diffuseRay = Geometry.sum(normal, Geometry.product(direction,Math.sin(Math.pow((rand.nextDouble()%2)-1,intersected.getMaterial().specularPower) * Math.PI)));
						Geometry.normalize(diffuseRay);
						
						Geometry.add(diff, Geometry.componentProduct(trace(scene,intersection,diffuseRay,intersected,level + 1),intersected.getMaterial().diffuse));
					}
					diffuse = true;
					Geometry.add(color, Geometry.product(diff, 1.0/diffuseDensity));
					
				}
			} else {
				Geometry.add(color, scene.background);
			}
		} 
		
		return color;
	}
	
	private double[] shade(Scene scene, Primitive primitive, double[] pos, double[] normal, double[] reflection) {
        double[] color = {0,0,0};
        
        if (ambientShading) {
                color = Geometry.componentProduct(primitive.getMaterial().ambient, scene.ambient);
        }
        
        double[] posToLight; // vector pointing from position to light source           
        
        double posToLightDistance; // distance between position and light source
        Light[] lights = scene.lights; // all lights in the scene
        double diffuse; // diffuse shading
        double specular; // specular shading
        
        double cosa; // angle between normal and ray to light
        
        double[] shade = {0,0,0};
        Primitive intersected = null; // primitive that intersects connection between position and light source (throws a shadow)
        
        // do shading (phong model)
        for (int i = 0; i < lights.length; i++) {
                // calculate vector between position and light source (normalized) and distance between them
                posToLight = Geometry.difference(lights[i].position, pos);
                posToLightDistance = Geometry.norm(posToLight);
                Geometry.multiply(posToLight, 1/posToLightDistance);

                // angle between normal and vector pointing from position to light
                cosa = (Geometry.dotProduct(normal, posToLight));
                
                if (cosa > 0) { // light hits surface on the front side                        
                        
                        // check for intersections between position and light source
                        if (shadowProbe) {
                                intersected = intersection(scene, primitive, pos, posToLight, posToLightDistance, null);                                
                                if (softShadows) {
                                	diffuse = 0;
                                	
            						// send random rays to light source
            						double[] random, direction, softRay;
            						double phiMax = Math.atan(lights[i].radius / posToLightDistance); // maximal angle
            						Primitive softIntersected;
            						
            						int density = softShadowDensity;
                                	
                                	for (int j = 0; j < density; j++) {
            							random = new double[]{rand.nextDouble()%1,rand.nextDouble()%1,rand.nextDouble()%1}; // random ray
            							
            							direction = Geometry.crossProduct(posToLight, random); // direction of soft ray
            							Geometry.normalize(direction);
            							
            							Geometry.multiply(direction, Math.sin( phiMax*((rand.nextDouble()%2)-1) ) );            							
            							
            							softRay = Geometry.sum(posToLight, direction);
            							Geometry.normalize(softRay);
            							
            							softIntersected = intersection(scene, primitive, pos, softRay, posToLightDistance, null);
            							
            							if (softIntersected == null) {
            								diffuse += (cosa * lights[i].brightness);
            							}
										
									}
                                	
                                	diffuse /= density;
                                	
                                } else { // no soft shadows
	                                if (intersected == null) {
	                                	diffuse = cosa * lights[i].brightness;
	                                } else {
	                                	diffuse = cosa * lights[i].brightness * intersected.getMaterial().translucence;
	                                }
                                }                                
                        } else { // no shadow probe
                        	diffuse = cosa * lights[i].brightness;
                        }
                        
                        // calculate diffuse reflection
                        if (diffuseShading) {
                            shade = Geometry.product(primitive.getMaterial().diffuse, diffuse);
                            shade = Geometry.componentProduct(shade, lights[i].color);
                            // apply diffuse reflection
                            Geometry.add(color, shade);
                        }
                        
                        if (intersected == null) { // no other primitive is blocking light from light source
                                // calculate specular reflection (angle between reflected ray vector pointing from position to light)
                                if (specularShading) {
                                        specular = Geometry.dotProduct(reflection, posToLight);
                                        if (specular > 0) {
                                                specular = Math.pow(specular, primitive.getMaterial().specularPower);
                                                specular = specular * lights[i].brightness;
                                                shade = Geometry.product(primitive.getMaterial().specular, specular);
                                                shade = Geometry.componentProduct(shade, lights[i].color);
                                                // apply specular reflection
                                                Geometry.add(color, shade);
                                        }
                                }
                        }
                }
        }
        
        return color;
	}
	
	public void save(String fname) {		
		File file = new File(fname+".png");
		try {
			ImageIO.write(this, "png", file);
			System.out.println("saved image to " + fname+ ".png");
		} catch (IOException e) {
			e.printStackTrace();
		}
		
	}



}
